package com.akto.action.gpt;

import com.akto.action.UserAction;
import com.akto.action.gpt.cache.VulnerabilityAnalysisCache;
import com.akto.action.gpt.handlers.GptQuery;
import com.akto.action.gpt.handlers.QueryHandler;
import com.akto.action.gpt.handlers.QueryHandlerFactory;
import com.akto.dto.User;
import com.mongodb.BasicDBObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Generic action class for content analysis using GPT.
 * Supports different analysis types like RedTeaming, Threat and Guardrail Activity, etc.
 */
public class VulnerabilityAnalysisAction extends UserAction {
    
    private static final Logger logger = LoggerFactory.getLogger(VulnerabilityAnalysisAction.class);
    public static final String USER_EMAIL = "user_email";
    
    private final VulnerabilityAnalysisCache cache = VulnerabilityAnalysisCache.getInstance();
    
    // Generic input parameters
    private String requestData;        // JSON string containing the data to analyze
    private String analysisType;       // Type of analysis: "redteaming", "threat", "guardrail"
    private String resultId;           // Optional ID for caching (e.g., testResultId, activityId, reportId)
    
    
    // Generic output
    private BasicDBObject analysisResult;
    
    public String analyzeVulnerability() {
        try {
            User sUser = getSUser();
            
            // Validate required parameters
            if (requestData == null || requestData.trim().isEmpty()) {
                throw new IllegalArgumentException("Request data is required");
            }
            
            if (analysisType == null || analysisType.trim().isEmpty()) {
                throw new IllegalArgumentException("Analysis type is required");
            }
            
            // Parse the structured data from frontend
            BasicDBObject contentData = parseContentData(requestData);
            
            // Extract result ID for caching (generic for any analysis type)
            String extractedResultId = contentData.getString("resultId");
            if (extractedResultId == null && resultId != null) {
                extractedResultId = resultId;
            }
            // Extract test context for metadata
            String testCategory = "UNKNOWN";
            String testDescription = "";
            BasicDBObject testContext = (BasicDBObject) contentData.get("testContext");
            if (testContext != null) {
                if (testContext.getString("category") != null) {
                    testCategory = testContext.getString("category");
                }
                if (testContext.getString("description") != null) {
                    testDescription = testContext.getString("description");
                }
            }
            
            // Check cache if result ID is available
            String cacheKey = cache.generateCacheKey(analysisType, extractedResultId, contentData);
            if (cacheKey != null) {
                analysisResult = cache.get(cacheKey);
                if (analysisResult != null) {
                    return SUCCESS.toUpperCase();
                }
            }
            
            // Prepare metadata for the query handler
            BasicDBObject meta = new BasicDBObject();
            meta.put(USER_EMAIL, sUser.getLogin());
            meta.put("contentData", contentData);
            meta.put("testCategory", testCategory);
            meta.put("testDescription", testDescription);
            
            // Get the appropriate query handler based on analysis type
            GptQuery query = getGptQueryForAnalysisType(analysisType);
            QueryHandler queryHandler = QueryHandlerFactory.getQueryHandler(query);
            analysisResult = queryHandler.handleQuery(meta);
            
            // Cache the result if we have a valid cache key
            if (analysisResult != null && cacheKey != null) {
                cache.put(cacheKey, analysisResult);
            }
            
            return SUCCESS.toUpperCase();
        } catch (Exception e) {
            logger.error("Error analyzing content for type: " + analysisType, e);
            addActionError(e.getMessage());
            return ERROR.toUpperCase();
        }
    }
    
    /**
     * Maps analysis type to the corresponding GptQuery enum
     */
    private GptQuery getGptQueryForAnalysisType(String analysisType) {
        switch (analysisType.toLowerCase()) {
            case "redteaming":
                return GptQuery.ANALYZE_VULNERABILITY;
            default:
                logger.warn("Unknown analysis type: {}, defaulting to redteaming", analysisType);
                return GptQuery.ANALYZE_VULNERABILITY;
        }
    }
    
    /**
     * Parses the input data string into a structured BasicDBObject
     */
    private BasicDBObject parseContentData(String requestData) {
        try {
            return BasicDBObject.parse(requestData);
        } catch (Exception e) {
            logger.warn("Failed to parse structured content data, using fallback", e);
            // Fallback to simple format if parsing fails
            BasicDBObject fallback = new BasicDBObject();
            fallback.put("response", new BasicDBObject("body", requestData));
            return fallback;
        }
    }
    
    
    // Getters and Setters
    public String getRequestData() {
        return requestData;
    }
    
    public void setRequestData(String requestData) {
        this.requestData = requestData;
    }
    
    public String getAnalysisType() {
        return analysisType;
    }
    
    public void setAnalysisType(String analysisType) {
        this.analysisType = analysisType;
    }
    
    public String getResultId() {
        return resultId;
    }
    
    public void setResultId(String resultId) {
        this.resultId = resultId;
    }
    
    public BasicDBObject getAnalysisResult() {
        return analysisResult;
    }
    
    public void setAnalysisResult(BasicDBObject analysisResult) {
        this.analysisResult = analysisResult;
    }
}
